import os
import pandas as pd
from PIL import Image
import io


PNG_SIGNATURE = b"\x89PNG"

def convert_parquet_images(
    input_dir: str,
    output_dir,
    image_column: str = 'image_data'
) -> None:
    """
    遍历 input_dir 下所有子文件夹中的 .parquet 文件，
    将指定列的二进制图像数据读取并保存为图片。

    参数:
        input_dir: 输入根目录，包含若干子文件夹
        output_dir: 图片保存根目录，结构与 input_dir 保持一致
        image_column: 存储图像二进制数据的列名，默认为 'image_data'
    """
    # 创建输出根目录
    if output_dir is None:
        output_dir = input_dir + '_raw'
    os.makedirs(output_dir, exist_ok=True)

    # 遍历目录
    for root, _, files in os.walk(input_dir):
        for filename in files:
            if 'Evol-Instruct-GPT4-Turbo' in root:
                continue
            
            if not filename.endswith('.parquet'):
                continue
            parquet_path = os.path.join(root, filename)
            # 读取 parquet
            df = pd.read_parquet(parquet_path)

            print(f"File: {parquet_path}\nColumns: {list(df.columns)}\n")
            
            # 相对路径，用于在输出目录中重建层级
            rel_dir = os.path.relpath(root, input_dir)
            out_subdir = os.path.join(output_dir, rel_dir)
            os.makedirs(out_subdir, exist_ok=True)

            # 遍历每一行，读取图像数据并保存
            for idx, row in df.iterrows():
                img_data = row[image_column]
                # print(img_data)
                # 如果存储为字节，要用 BytesIO
                if isinstance(img_data, (bytes, bytearray)):
                    img = Image.open(io.BytesIO(img_data))
                elif isinstance(img_data, dict):
                    img = None
                    for v in img_data.values():
                        if isinstance(v, (bytes, bytearray)) and v.startswith(PNG_SIGNATURE):
                            img = Image.open(io.BytesIO(v))
                            break
                    if img is None:
                        print(f"Row {idx} dict contains no valid PNG bytes, skip.")
                        continue
                else:
                    # 如果是 NumPy 数组
                    img = Image.fromarray(img_data)
                    # img = img_data

                # 构造输出文件名：原 parquet 名 + 行号
                base_name = os.path.splitext(filename)[0]
                out_name = f"{base_name}_{idx}.png"
                out_path = os.path.join(out_subdir, out_name)
                img.save(out_path)
                print(f"Saved image: {out_path}")

if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser(
        description='Convert all parquet-stored images under a directory to PNG files'
    )
    parser.add_argument('input_dir', help='根目录，包含子文件夹中的 parquet 文件')
    parser.add_argument('--output_dir')
    parser.add_argument(
        '--column',
        default='image',
        help='Parquet 中保存二进制图像数据的列名，默认 image_data'
    )
    args = parser.parse_args()

    convert_parquet_images(
        input_dir=args.input_dir,
        output_dir=args.output_dir,
        image_column=args.column
    )
